# -*- coding: utf-8 -*-
"""
Code Introduction:
    This code 
Version History:
    Created: Tue Jan 12 14:21:49 2021
    Current: 

@author: Xing Guo (guoxing.econ@gmail.com)

"""

#%% Import Moduels

## Data Process Tools
import pandas as pd
import datetime

## Numerical
import numpy as np

## Statistical
import statsmodels.formula.api as sm
import statsmodels.api as SMAPI

# End of Section: Import Moduels
###############################################################################

### Local Projection Functions
## For Testing
# LP = LP_Info('InterestRate')
# DS_MS = pickle.load(open('TempData\DS_MonetaryPolicy.p','rb'))
# DS_DepVar = pickle.load(open('TempData\DS_DepVar_National.p','rb'))
# DS = DS_MS.join(DS_DepVar).sort_index().reset_index()

#%% Local Projection Regression


### Information Specifying Local Projection Regression
class LP_Info:
    def __init__(self,DepVar,ShockVar='MS',ControlVarList=[], \
                      DateVar='QDate', \
                      SamplePeriod=[datetime.datetime(1974,1,1),datetime.datetime(2015,1,1)], \
                      MaxLag_Shock=4,MaxLag_Control=8,MaxH=20,ConfIntAlpha=0.1):
        self.DepVar = DepVar
        self.ShockVar = ShockVar
        self.ControlVarList = ControlVarList
        self.DateVar = DateVar
        self.SamplePeriod = SamplePeriod
        self.MaxLag_Shock = MaxLag_Shock
        self.MaxLag_Control = MaxLag_Control
        self.MaxH = MaxH
        self.ConfIntAlpha = ConfIntAlpha


### Local Projection Regression
def LP_Reg(DS,LP):
    RegDS = DS.copy()
    ## Compose the List of Control Variables
    RegControlVarList = []
    # Lagged Monetary Shocks
    if LP.MaxLag_Shock>0:
        for ll in range(1,LP.MaxLag_Shock+1):
            TempVar = LP.ShockVar+'_Lag_'+str(ll)
            RegDS[TempVar] = RegDS[LP.ShockVar].shift(ll)
            RegControlVarList.append(TempVar)
    # Lagged Control Variables
    if len(LP.ControlVarList)>0 and LP.MaxLag_Control>0:
        for vv in LP.ControlVarList:
            for ll in range(1,LP.MaxLag_Control+1):
                TempVar = vv+'_Lag_'+str(ll)
                RegDS[TempVar] = RegDS[vv].shift(ll)
                RegControlVarList.append(TempVar)
    
    ## Unit Regression
    def LocalProjection_Unit(h):
         RegDS['YY'] = RegDS[LP.DepVar].shift(-h)-RegDS[LP.DepVar].shift(1)
         RegDesign = 'YY'+'~'+LP.ShockVar+'+'+'+'.join(RegControlVarList)
         TempInd = (RegDS[LP.DateVar]>=LP.SamplePeriod[0]) & (RegDS[LP.DateVar]<=LP.SamplePeriod[1])
         RegModel = sm.ols(formula=RegDesign,data=RegDS[TempInd]).fit()
         ConfInt = RegModel.conf_int(alpha=LP.ConfIntAlpha)
         return {'Coef': RegModel.params[LP.ShockVar], \
                 'Std': RegModel.bse[LP.ShockVar], \
                 'Tvalue':RegModel.tvalues[LP.ShockVar], \
                 'CI_Left':ConfInt.loc[LP.ShockVar,0], \
                 'CI_Right':ConfInt.loc[LP.ShockVar,1] }
    ## Collect the Regression Results
    RegResult_DictList = [LocalProjection_Unit(hh) for hh in range(LP.MaxH+1)]
    RegResult_DS = pd.concat([pd.Series(RR) for RR in RegResult_DictList], \
                             axis=1,names=list(range(LP.MaxH+1))).T
    
    return RegResult_DS

#%% Romer and Romer Regression

### Information Specifying Romer and Romer Regression
class RR_Info:
    def __init__(self,DepVar,ShockVar='MS',DateVar='QDate',IRFType='NonAccumulated', \
                      NonLagControlVarList=[],LagControlVarList=[], \
                      SamplePeriod=[datetime.datetime(1974,1,1),datetime.datetime(2015,1,1)], \
                      MinLag_Shock=0,MaxLag_Shock=20,MaxLag_Control=8,ConfIntAlpha=0.1):
        self.DepVar = DepVar
        self.ShockVar = ShockVar
        self.LagControlVarList = LagControlVarList
        self.NonLagControlVarList = NonLagControlVarList
        self.DateVar = DateVar
        self.SamplePeriod = SamplePeriod
        self.MinLag_Shock = MinLag_Shock
        self.MaxLag_Shock = MaxLag_Shock
        self.MaxLag_Control = MaxLag_Control
        self.IRFType = IRFType
        self.ConfIntAlpha = ConfIntAlpha


### Regression
def RR_Reg(DS,RR):
    RegDS = DS.copy()
    ## 1. Variable Lists
    # Non-Lagged Control
    NonLagControlVarList = RR.NonLagControlVarList
    # Lagged Control
    LagControlVarList = []
    if not RR.LagControlVarList:
        for ii in range(1,RR.MaxLag_Control+1):
            for vv in RR.LagControlVarList:
                TempVar = vv+'_Lag_'+str(ii)
                RegDS[TempVar] = RegDS[vv].shift(ii)
                LagControlVarList.append(TempVar)
    # Shocks
    ShockList = []
    for ii in range(RR.MinLag_Shock,RR.MaxLag_Shock+1):
        TempVar = RR.ShockVar+'_Lag_'+str(ii)
        RegDS[TempVar] = RegDS[RR.ShockVar].shift(ii)
        ShockList.append(TempVar)
    ## 2. Regression
    XVarList = ShockList+NonLagControlVarList+LagControlVarList
    RegDesign = RR.DepVar+'~'+'+'.join(XVarList)
    TempInd = (RegDS[RR.DateVar]>=RR.SamplePeriod[0]) & (RegDS[RR.DateVar]<=RR.SamplePeriod[1])
    RegFit = sm.ols(formula=RegDesign,data=RegDS[TempInd]).fit()
    ## 3. IRF: Point Estimates and Std/CI
    TempDict = {RR.ShockVar+'_Lag_'+str(ii): ii for ii in range(RR.MinLag_Shock,RR.MaxLag_Shock+1)}
    Coef = RegFit.params[ShockList].rename(index=TempDict)
    if RR.IRFType=='Accumulated':
        Coef = Coef.cumsum()
        
        TempFullCov = RegFit.cov_params().loc[ShockList,ShockList] \
                      .rename(index=TempDict,columns=TempDict)
        
        TempFullCov.sort_index(axis=0,inplace=True)
        TempFullCov.sort_index(axis=1,inplace=True)
        Std = []
        TempLagList = list(TempFullCov.index)
        for ii in range(len(TempLagList)):
            TempStd = np.sqrt( np.ones([1,ii+1])@ \
                               TempFullCov.loc[TempLagList[0:ii+1],TempLagList[0:ii+1]]@ \
                               np.ones([ii+1,1]) )[0][0]
            Std.append(TempStd)
        Std = pd.Series(Std,index=TempLagList)
        
        RegResult = pd.concat([Coef,Std],axis=1,keys=['Coef','Std'])
    else:
        Std = RegFit.bse[ShockList].rename(index=TempDict)
        Tvalue = RegFit.tvalues[ShockList].rename(index=TempDict)
        ConfInt = RegFit.conf_int(alpha=RR.ConfIntAlpha)
        CI_Left = ConfInt.loc[ShockList,0].rename(index=TempDict)
        CI_Right = ConfInt.loc[ShockList,1].rename(index=TempDict)
        
        RegResult = pd.concat([Coef,Std,Tvalue,CI_Left,CI_Right],axis=1, \
                              keys=['Coef','Std','Tvalue','CI_Left','CI_Right'])
    
 
    return RegResult
        